from keras import Model
from keras.models import load_model
import model
from model import Nets_2D
import matplotlib.pyplot as plt
import numpy as np
import cv2
import time

# 仅支持单数字
prefix = 'jpg'
index = 3         # 0-9总共十张测试图片
test = plt.imread('test_image/{}_test.{}'.format(index,prefix))
# test = cv2.cvtColor(test,cv2.COLOR_RGB2GRAY)
# test = cv2.resize(test,dsize=(28,28))

plt.imshow(test,cmap='gray')
test = test[np.newaxis,:,:,np.newaxis]

# 参数设定
# width     = 28
# height    = 28
epochs    = 100
key       = 'mynet_1'

predict_methods = {
    'mynet_1' : 'checkpoints/mynet_1.h5',
    'mynet_2' : 'checkpoints/mynet_2.h5',
    'mynet_1_earlystop' : 'checkpoints/mynet_1_{}epochs.h5'.format(epochs),
    'mynet_2_earlystop' : 'checkpoints/mynet_2_{}epochs.h5'.format(epochs),
}

file_string = predict_methods[key]
m = load_model(file_string)
m.load_weights(file_string)

start = time.time()
result = m.predict(test,batch_size=1,verbose=1)
end = time.time()
classfication = np.argmax(result)
print('===========================')
print('<netout:{}>'.format(result))
print('<class:{}>'.format(classfication))
print('<The Time:%.2f>'%(end-start))
print('===========================')
plt.axis('off')
plt.show()

